package org.hong.monkey.network.protocol;

import com.google.common.base.Preconditions;
import io.netty.buffer.ByteBuf;
import io.netty.channel.FileRegion;
import io.netty.util.AbstractReferenceCounted;
import io.netty.util.ReferenceCountUtil;

import java.io.IOException;
import java.nio.channels.WritableByteChannel;

public class MessageWithHeader extends AbstractReferenceCounted implements FileRegion {

    private final ByteBuf header;
    private final int headerLength;
    private final Object body;
    private final long bodyLength;
    private long totalBytesTransferred;

    MessageWithHeader(ByteBuf header, Object body, long bodyLength) {
        Preconditions.checkArgument(body instanceof ByteBuf || body instanceof FileRegion,
                "Body must be a ByteBuf or a FileRegion.");
        this.header = header;
        this.headerLength = header.readableBytes();
        this.body = body;
        this.bodyLength = bodyLength;
    }

    @Override
    public long count() {
        return headerLength + bodyLength;
    }

    @Override
    public long position() {
        return 0;
    }

    @Override
    public long transfered() {
        return totalBytesTransferred;
    }

    @Override
    public long transferTo(final WritableByteChannel target, final long position) throws IOException {
        Preconditions.checkArgument(position == totalBytesTransferred, "Invalid position.");
        long writtenHeader = 0;
        if (header.readableBytes() > 0) {
            writtenHeader = copyByteBuf(header, target);
            totalBytesTransferred += writtenHeader;
            if (header.readableBytes() > 0) {
                return writtenHeader;
            }
        }

        // Bytes written for body in this call.
        long writtenBody = 0;
        if (body instanceof FileRegion) {
            writtenBody = ((FileRegion) body).transferTo(target, totalBytesTransferred - headerLength);
        } else if (body instanceof ByteBuf) {
            writtenBody = copyByteBuf((ByteBuf) body, target);
        }
        totalBytesTransferred += writtenBody;

        return writtenHeader + writtenBody;
    }

    @Override
    protected void deallocate() {
        header.release();
        ReferenceCountUtil.release(body);
    }

    private int copyByteBuf(ByteBuf buf, WritableByteChannel target) throws IOException {
        int written = target.write(buf.nioBuffer());
        buf.skipBytes(written);
        return written;
    }
}
